//by mircemk, June 2026

#include "unihiker_k10.h"
#include <TFT_eSPI.h>
#include "esp_camera.h"
#include <Preferences.h>

#define Serial0 Serial

UNIHIKER_K10 k10;
TFT_eSPI tft = TFT_eSPI();
Preferences prefs;

uint8_t screen_dir = 2;

#define XCLK_GPIO_NUM  7
#define SIOD_GPIO_NUM  47
#define SIOC_GPIO_NUM  48
#define Y9_GPIO_NUM    6
#define Y8_GPIO_NUM    15
#define Y7_GPIO_NUM    16
#define Y6_GPIO_NUM    18
#define Y5_GPIO_NUM    9
#define Y4_GPIO_NUM    11
#define Y3_GPIO_NUM    10
#define Y2_GPIO_NUM    8
#define VSYNC_GPIO_NUM 4
#define HREF_GPIO_NUM  5
#define PCLK_GPIO_NUM  17

#define FEAT_SIZE 64
#define NUM_CLASSES 7

uint8_t currentFeat[FEAT_SIZE];
uint8_t refs[NUM_CLASSES][FEAT_SIZE];
bool hasRef[NUM_CLASSES];

const char commands[NUM_CLASSES] = {'e', 'c', 't', 'i', 's', 'b', 'r'};

const char* labels[NUM_CLASSES] = {
  "EMPTY",
  "CAPACITOR",
  "TRANSISTOR",
  "IC CHIP",
  "STM32",
  "BATTERY",
  "RESISTOR"
};

const char* keys[NUM_CLASSES] = {
  "empty",
  "cap",
  "trans",
  "ic",
  "stm32",
  "bat",
  "res"
};

void drawStatus(String text, uint16_t color) {
  tft.fillRect(0, 280, 240, 40, TFT_BLACK);
  tft.drawRect(0, 280, 240, 40, TFT_YELLOW);
  tft.setTextColor(color, TFT_BLACK);
  tft.setTextSize(2);
  tft.drawString(text, 8, 292);
}

void saveReference(int index) {
  prefs.putBytes(keys[index], refs[index], FEAT_SIZE);
}

bool loadReference(int index) {
  size_t len = prefs.getBytesLength(keys[index]);
  if (len != FEAT_SIZE) return false;
  prefs.getBytes(keys[index], refs[index], FEAT_SIZE);
  return true;
}

void setupCamera() {
  camera_config_t config;

  config.ledc_channel = LEDC_CHANNEL_0;
  config.ledc_timer   = LEDC_TIMER_0;

  config.pin_d0 = Y2_GPIO_NUM;
  config.pin_d1 = Y3_GPIO_NUM;
  config.pin_d2 = Y4_GPIO_NUM;
  config.pin_d3 = Y5_GPIO_NUM;
  config.pin_d4 = Y6_GPIO_NUM;
  config.pin_d5 = Y7_GPIO_NUM;
  config.pin_d6 = Y8_GPIO_NUM;
  config.pin_d7 = Y9_GPIO_NUM;

  config.pin_xclk  = XCLK_GPIO_NUM;
  config.pin_pclk  = PCLK_GPIO_NUM;
  config.pin_vsync = VSYNC_GPIO_NUM;
  config.pin_href  = HREF_GPIO_NUM;

  config.pin_sscb_sda = SIOD_GPIO_NUM;
  config.pin_sscb_scl = SIOC_GPIO_NUM;

  config.pin_pwdn  = -1;
  config.pin_reset = -1;

  config.xclk_freq_hz = 20000000;
  config.pixel_format = PIXFORMAT_RGB565;
  config.frame_size   = FRAMESIZE_QVGA;

  config.fb_count = 1;
  config.fb_location = CAMERA_FB_IN_PSRAM;
  config.grab_mode = CAMERA_GRAB_LATEST;

  if (esp_camera_init(&config) != ESP_OK) {
    drawStatus("CAM ERROR", TFT_RED);
    while (1) delay(1000);
  }

  sensor_t *s = esp_camera_sensor_get();
  s->set_vflip(s, 1);
  s->set_hmirror(s, 0);
}

void extractFeatures(camera_fb_t *fb, uint8_t *feat) {
  uint16_t *pix = (uint16_t *)fb->buf;

  int roiX = 40;
  int roiY = 60;
  int roiW = 160;
  int roiH = 180;

  int index = 0;

  for (int gy = 0; gy < 8; gy++) {
    for (int gx = 0; gx < 8; gx++) {
      long sum = 0;
      int count = 0;

      int x0 = roiX + gx * roiW / 8;
      int x1 = roiX + (gx + 1) * roiW / 8;
      int y0 = roiY + gy * roiH / 8;
      int y1 = roiY + (gy + 1) * roiH / 8;

      for (int y = y0; y < y1; y += 3) {
        for (int x = x0; x < x1; x += 3) {
          uint16_t c = pix[y * fb->width + x];

          uint8_t r = ((c >> 11) & 0x1F) << 3;
          uint8_t g = ((c >> 5)  & 0x3F) << 2;
          uint8_t b = (c & 0x1F) << 3;

          uint8_t gray = (r + g + b) / 3;
          sum += gray;
          count++;
        }
      }

      feat[index++] = sum / count;
    }
  }
}

int distanceTo(uint8_t *a, uint8_t *b) {
  long d = 0;

  for (int i = 0; i < FEAT_SIZE; i++) {
    int diff = (int)a[i] - (int)b[i];
    if (diff < 0) diff = -diff;
    d += diff;
  }

  return d / FEAT_SIZE;
}

int commandToIndex(char cmd) {
  for (int i = 0; i < NUM_CLASSES; i++) {
    if (cmd == commands[i]) return i;
  }
  return -1;
}

String classifyObject() {
  int best = 9999;
  int bestIndex = -1;

  for (int i = 0; i < NUM_CLASSES; i++) {
    if (!hasRef[i]) continue;

    int d = distanceTo(currentFeat, refs[i]);

    if (d < best) {
      best = d;
      bestIndex = i;
    }
  }

  Serial0.print("BEST distance=");
  Serial0.println(best);

  if (bestIndex < 0) return "NO REFS";
  if (best > 25) return "UNKNOWN";

  return String(labels[bestIndex]);
}

void learnReference(char cmd) {
  int index = commandToIndex(cmd);

  if (index < 0) return;

  memcpy(refs[index], currentFeat, FEAT_SIZE);
  hasRef[index] = true;
  saveReference(index);

  String msg = "SAVED ";
  msg += labels[index];

  drawStatus(msg, TFT_GREEN);

  Serial0.print("Learned and saved: ");
  Serial0.println(labels[index]);
}

void setup() {
  Serial0.begin(115200);
  delay(1000);

  k10.begin();
  k10.initScreen(screen_dir);

  tft.init();
  tft.setRotation(2);
  tft.setSwapBytes(false);
  tft.fillScreen(TFT_BLACK);

  prefs.begin("refs7", false);

  int loaded = 0;
  for (int i = 0; i < NUM_CLASSES; i++) {
    hasRef[i] = loadReference(i);
    if (hasRef[i]) loaded++;
  }

  setupCamera();

  String startMsg = "LOADED ";
  startMsg += String(loaded);
  startMsg += "/7";
  drawStatus(startMsg, TFT_GREEN);

  Serial0.println("Commands:");
  Serial0.println("e = learn EMPTY");
  Serial0.println("c = learn CAPACITOR");
  Serial0.println("t = learn TRANSISTOR");
  Serial0.println("i = learn IC CHIP");
  Serial0.println("s = learn STM32");
  Serial0.println("b = learn BATTERY");
  Serial0.println("r = learn RESISTOR");
}

void loop() {
  camera_fb_t *fb = esp_camera_fb_get();

  if (!fb) {
    delay(200);
    return;
  }

  tft.pushImage(0, 0, 240, 280, (uint16_t *)fb->buf);

  extractFeatures(fb, currentFeat);

  esp_camera_fb_return(fb);

  if (Serial0.available()) {
    char cmd = Serial0.read();

    if (commandToIndex(cmd) >= 0) {
      learnReference(cmd);
      delay(1000);
      return;
    }
  }

  static unsigned long lastClassify = 0;

  if (millis() - lastClassify > 700) {
    lastClassify = millis();

    String result = classifyObject();

    if (result == "UNKNOWN" || result == "NO REFS") {
      drawStatus(result, TFT_RED);
    } else if (result == "EMPTY") {
      drawStatus(result, TFT_YELLOW);
    } else {
      drawStatus(result, TFT_GREEN);
    }
  }

  delay(80);
}